package com.agentframework.core.agent;

import com.agentframework.core.action.Action;
import com.agentframework.core.environment.AbstractEnvironment;
import org.springframework.ai.chat.messages.Message;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * @author wangjun
 **/
public class AgentContext {

    private AbstractEnvironment environment;

    private Class<? extends ReactStrategy> reactMode;

    private final BlockingQueue<Message> messageDeque;

    /**
     * 异常处理消息
     * 比如大模型没有按规定的Json结构返回，导致解析json失败等异常
     */
    private final ConcurrentLinkedQueue<Message> errorMessageQueue = new ConcurrentLinkedQueue<>();

    private final Set<String> watch = new HashSet<>();

    private final List<Message> news = new ArrayList<>();

    private int maxReactLoop = 1;

    private Action todo;

    /**
     * -1 表示初始或者是终止状体，其中todo为空，返回具体action的index值
     */
    private int actionIndex = -1;

    public AgentContext(Class<? extends ReactStrategy> agentReactMode) {
        this.reactMode = agentReactMode;
        this.messageDeque = new LinkedBlockingQueue<>(1000);
    }

    public AgentContext() {
        this(InOrderReactStrategy.class);
    }

    public void setEnvironment(AbstractEnvironment environment) {
        this.environment = environment;
    }

    public AbstractEnvironment getEnvironment() {
        return environment;
    }

    public void putMessage(Message message) {
        try {
            messageDeque.put(message);
        } catch (InterruptedException e) {
            throw new RuntimeException("消息投递失败", e);
        }
    }

    public List<Message> popAllMessage() {
        return new ArrayList<>(messageDeque);
    }

    public Message takeMessage() throws InterruptedException {
        return messageDeque.take();
    }

    public int getMessageSize() {
        return messageDeque.size();
    }

    public Action getTodo() {
        return todo;
    }

    public void setTodo(Action todo) {
        this.todo = todo;
    }

    public void setReactMode(Class<? extends ReactStrategy> reactMode) {
        this.reactMode = reactMode;
    }

    public Class<? extends ReactStrategy> getReactMode() {
        return reactMode;
    }

    public int getMaxReactLoop() {
        return maxReactLoop;
    }

    public void setMaxReactLoop(int maxReactLoop) {
        this.maxReactLoop = maxReactLoop;
    }

    public int getActionIndex() {
        return actionIndex;
    }

    public void setActionIndex(int index) {
        this.actionIndex = index;
    }

    public List<Message> getNews() {
        return news;
    }

    public Set<String> getWatch() {
        return watch;
    }

    public void addWatch(String agentId) {
        getWatch().add(agentId);
    }

    public void addWatches(List<String> agentIds) {
        getWatch().addAll(agentIds);
    }

    public void addErrorMessage(Message message) {
        errorMessageQueue.offer(message);
    }

    public Message getErrorMessage() {
        return errorMessageQueue.poll();
    }

}
